import os
import cv2
dic_k2y = {
        'Car': 0, 
        'Van': 1, 
        'Truck': 2, 
        'Tram': 3, 
        'Pedestrian': 4, 
        'Person_sitting': 5, 
        'Cyclist': 6, 
        'Misc': 7, 
        }
dic_y2k = { value: key for key, value in  dic_k2y.items()}


def kitti_2_yolo(img_path, label_path, save_path = '../converted_label/k_2_y'):
    dic = dic_k2y
    try: 
        os.makedirs(save_path)
    except:
        None
    for done_num, img_file in enumerate(os.listdir(img_path)):
        if done_num % 1000 == 0:
            print(done_num)
        if img_file.endswith(('.jpg', '.png', '.jpeg')):  # 确保是图片文件
            # 搞到图片和对应的label路径
            label_file = os.path.join(label_path, os.path.splitext(img_file)[0] + ".txt")
            save_path_this = os.path.join(save_path, os.path.splitext(img_file)[0] + ".txt")
            img_file = os.path.join(img_path, img_file)
            #print(save_path_this)

            #读取图片的size
            image = cv2.imread(img_file)
            h, w, c = image.shape
            dh = 1.0 / h
            dw = 1.0 / w
            

            #读label
            labels = []#一张图片的label可能不止一行
            with open(label_file,"r") as f:
                for line in f.readlines():
                    labels.append(line.strip('\n').split())

            with open(save_path_this, 'w+') as write_:
                for label in labels:
                    #具体到每个框的label
                    if label[0] == 'DontCare':
                        continue
                    type_id = dic[label[0]]
                    kitti_lx = float(label[4])
                    kitti_ly = float(label[5])
                    kitti_rx = float(label[6])
                    kitti_ry = float(label[7])
    
                    w = (kitti_rx - kitti_lx) * dw
                    h = (kitti_ry - kitti_ly) * dh
                    x = (kitti_rx + kitti_lx) / 2.0 * dw
                    y = (kitti_ry + kitti_ly) / 2.0 * dh
    
                    # 通常保留六位小鼠
                    w = round(w, 6)
                    h = round(h, 6)
                    x = round(x, 6)
                    y = round(y, 6)
                    
                    write_.write(f'{type_id} {x} {y} {w} {h}\n')
    

def yolo_2_kitti(img_path, yolo_labels_path, save_path = '../converted_label/y_2_k'):
    try: 
        os.makedirs(save_path)
    except:
        None
    labels = os.listdir(yolo_labels_path)
    for i, label in enumerate(labels):
        if i % 1000 == 0:
            print(f'{i}/{len(labels)} has been done')
        if label.endswith(('.txt')):  # 确保是txt文件
            label_file = os.path.join(yolo_labels_path, label)
            save_path_this = os.path.join(save_path, label)
            img_file = os.path.join(img_path, os.path.splitext(label)[0] + ".png")
            image = cv2.imread(img_file)
            o_h, o_w, o_c = image.shape
            #为了逆归一化
            
            lines = []
            with open(label_file,"r") as f:
                for line in f.readlines():
                    lines.append(line.strip('\n').split())

            with open(save_path_this, 'w+') as write_:
                for object_ in lines:
                    object_type =  dic_y2k[eval(object_[0])]
                    c_x = float(object_[1])
                    c_y = float(object_[2])
                    w = float(object_[3])
                    h = float(object_[4])
        
                    box_w = w * o_w
                    box_h = h * o_h
                    box_c_x = c_x * o_w
                    box_c_y = c_y * o_h
                    box_l_x = (box_c_x - box_w / 2)
                    box_l_y = (box_c_y - box_h / 2)
                    box_r_x = (box_c_x + box_w / 2)
                    box_r_y = (box_c_y + box_h / 2)

                    write_.write(f'{object_type} 0.00 0 0.00 {box_l_x:.2f} {box_l_y:.2f} {box_r_x:.2f} {box_r_y:.2f} 0.00 0.00 0.00 0.00 0.00 0.00 0.00\n')

                
                
                
                
                
                
                